import os
import json
import time
import traceback
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options


class NetworkRecorder:
    """记录浏览器网络请求的类"""
    
    def __init__(self, save_dir=None, chrome_driver_path=None):
        """
        初始化网络记录器
        
        参数:
            save_dir: 保存记录的目录，如果为None，则使用当前目录
            chrome_driver_path: Chrome驱动程序的路径，如果为None，将使用当前目录下的chromedriver.exe
        """
        if save_dir is None:
            self.save_dir = os.path.dirname(os.path.abspath(__file__))
        else:
            self.save_dir = save_dir
            if not os.path.exists(self.save_dir):
                os.makedirs(self.save_dir)
        
        # 设置Chrome驱动路径
        if chrome_driver_path is None:
            # 默认使用同目录下的chromedriver.exe
            self.chrome_driver_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'chromedriver.exe')
        else:
            self.chrome_driver_path = chrome_driver_path
        
        # 验证驱动文件是否存在
        if not os.path.exists(self.chrome_driver_path):
            raise FileNotFoundError(f"找不到Chrome驱动文件: {self.chrome_driver_path}")
        
        print(f"网络请求记录将保存到: {self.save_dir}")
        print(f"使用Chrome驱动: {self.chrome_driver_path}")
        
        # 设置Chrome浏览器选项
        self.options = Options()
        self.options.add_argument('--start-maximized')  # 最大化窗口
        self.options.add_argument('--disable-extensions')  # 禁用扩展
        self.options.add_argument('--disable-gpu')  # 禁用GPU加速
        
        # 设置日志首选项，启用性能日志
        self.options.set_capability('goog:loggingPrefs', {'performance': 'ALL'})
        
        self.driver = None
        self.recording = False
        self.processed_requests = set()  # 用于记录已处理的请求
    
    def start_browser(self, url):
        """
        启动浏览器并访问指定URL
        
        参数:
            url: 要访问的网址
        """
        try:
            print(f"正在启动Chrome浏览器...")
            service = Service(executable_path=self.chrome_driver_path)
            
            self.driver = webdriver.Chrome(
                service=service,
                options=self.options
            )
            
            # 启用网络域以获取网络事件
            self.driver.execute_cdp_cmd('Network.enable', {})
            
            self.driver.get(url)
            self.recording = True
            print(f"浏览器已启动并访问: {url}")
            print(f"请在浏览器中进行操作，程序将记录所有JSON响应的网络请求...")
            print(f"按Ctrl+C可停止记录并关闭浏览器")
            return True
        except Exception as e:
            print(f"启动浏览器时出错: {e}")
            traceback.print_exc()
            return False
    
    def record_network(self):
        """持续记录网络请求并保存JSON响应"""
        if not self.driver:
            print("浏览器未启动，无法记录网络请求")
            return
            
        try:
            print("开始记录网络请求...")
            while self.recording:
                # 获取性能日志
                logs = self.driver.get_log('performance')
                
                for log in logs:
                    try:
                        log_data = json.loads(log['message'])
                        message = log_data.get('message', {})
                        method = message.get('method', '')
                        
                        # 监听响应接收事件
                        if method == 'Network.responseReceived':
                            params = message.get('params', {})
                            request_id = params.get('requestId')
                            response = params.get('response', {})
                            
                            if request_id and request_id not in self.processed_requests:
                                # 检查MIME类型是否为JSON
                                mime_type = response.get('mimeType', '')
                                
                                if 'json' in mime_type.lower() or 'application/json' in mime_type.lower():
                                    try:
                                        # 获取响应体
                                        result = self.driver.execute_cdp_cmd('Network.getResponseBody', {'requestId': request_id})
                                        body = result.get('body', '')
                                        
                                        # 确保是有效的JSON
                                        try:
                                            json_data = json.loads(body)
                                            
                                            # 获取请求URL
                                            url = response.get('url', '')
                                            
                                            # 创建文件名
                                            timestamp = datetime.now().strftime('%Y%m%d%H%M%S%f')[:16]
                                            filename = f"{timestamp}.txt"
                                            filepath = os.path.join(self.save_dir, filename)
                                            
                                            # 获取请求头
                                            headers = response.get('headers', {})
                                            
                                            # 获取请求体（可能无法获取）
                                            request_body = ''
                                            try:
                                                request_info = self.driver.execute_cdp_cmd('Network.getRequestPostData', {'requestId': request_id})
                                                request_body = request_info.get('postData', '')
                                            except:
                                                pass  # 忽略错误
                                            
                                            # 保存网络记录
                                            with open(filepath, 'w', encoding='utf-8') as f:
                                                f.write(f"url: {url}\n\n")
                                                f.write("请求头和请求体：\n")
                                                f.write("```\n")
                                                request_content = {
                                                    "headers": headers,
                                                    "referrer": headers.get('Referer', self.driver.current_url),
                                                    "referrerPolicy": "strict-origin-when-cross-origin",
                                                    "body": request_body,
                                                    "method": params.get('type', 'GET'),
                                                    "mode": "cors",
                                                    "credentials": "include"
                                                }
                                                f.write(json.dumps(request_content, indent=2, ensure_ascii=False))
                                                f.write("\n```\n\n")
                                                
                                                f.write("响应内容：\n")
                                                f.write("```\n")
                                                f.write(json.dumps(json_data, indent=2, ensure_ascii=False))
                                                f.write("\n```\n")
                                            
                                            print(f"已保存JSON响应记录到: {filename}")
                                        except json.JSONDecodeError:
                                            # 不是有效的JSON，跳过
                                            pass
                                    except Exception as e:
                                        # 忽略无法获取响应体的请求
                                        pass
                                
                                # 标记为已处理
                                self.processed_requests.add(request_id)
                    except:
                        # 忽略单条日志处理错误
                        pass
                
                time.sleep(0.5)  # 短暂休息，减少CPU使用率
                
        except KeyboardInterrupt:
            print("\n用户停止记录")
        except Exception as e:
            print(f"记录网络请求时出错: {e}")
            traceback.print_exc()
        finally:
            self.stop_recording()
    
    def stop_recording(self):
        """停止记录并关闭浏览器"""
        self.recording = False
        if self.driver:
            try:
                self.driver.quit()
                print("浏览器已关闭")
            except:
                print("关闭浏览器时出错")
            finally:
                self.driver = None
        print("记录已停止")


def main():
    """主函数"""
    try:
        print("===== 浏览器网络请求记录器 =====")
        print("此工具用于记录Chrome浏览器中的网络请求，特别是JSON格式的响应")
        print("------------------------------")
        
        url = input("请输入要访问的网址 (默认为东方航空促销页): ").strip()
        if not url:
            url = "https://www.ceair.com/activity/on-sale-new?eyJjb2RlIjoiSkhZWDE4In0=encoded"
        
        save_dir = input("请输入保存记录的目录 (默认为当前目录): ").strip()
        if not save_dir:
            save_dir = None
        
        # 使用默认驱动路径（同目录下的chromedriver.exe）
        recorder = NetworkRecorder(save_dir)
        if recorder.start_browser(url):
            recorder.record_network()
        else:
            print("启动浏览器失败，程序退出")
    except Exception as e:
        print(f"程序运行出错: {e}")
        traceback.print_exc()


if __name__ == "__main__":
    main() 